package main

import (
	"errors"
	"fmt"
	"github.com/spf13/pflag"
	"ntrace/bootstrap"
	"ntrace/core"
	"ntrace/utils/nets"
	"ntrace/utils/routines"
	"os"
	"strconv"
	"strings"
	"time"
)

var help = pflag.BoolP("help", "h", false, "help info")
var tcp = pflag.BoolP("tcp", "T", false, "use tcp syn packet instead of icmp echo")
var udp = pflag.BoolP("udp", "U", false, "use udp datagram instead of icmp echo")
var sources = pflag.StringP("sources", "s", "", "source address, while use icmp echo, use multi sources to cover ecmp paths")
var srcPorts = pflag.String("sport", "63000-63999", "source port or port range, while use udp or tcp, use multi sport and dport to cover ecmp paths")
var dstPorts = pflag.String("dport", "53", "destination port or port range, while use udp or tcp, use multi sport and dport to cover ecmp paths")
var tos = pflag.Uint8("tos", 0, "tos")
var maxTTL = pflag.Uint8("max-ttl", 32, "max ttl")
var length = pflag.IntP("length", "l", 64, "payload data length")
var timeout = pflag.Int64P("timeout", "t", 1000, "message response timeout time, measured in milliseconds")
var flowInterval = pflag.Int64("flow-interval", 100, "flow interval time, measured in milliseconds")
var printInterval = pflag.Int64("print-interval", 5000, "interval time of print statistical data, measured in milliseconds")
var count = pflag.IntP("count", "c", 10, "repeat count of trace flow")

func main() {

	if helpFunc() {
		return
	}

	err := bootstrap.InitEnv()
	if err != nil {
		panic(err)
	}

	task := buildTraceTask()
	if task == nil {
		return
	}

	agent, err := core.NewAgent(task)
	if err != nil {
		return
	}

	err = agent.StartAgent()
	if err != nil {
		agent.StopAgent()
	}

	routines.WaitExit()
	agent.StopAgent()
	time.Sleep(time.Second)
}

func buildTraceTask() *core.TraceTask {

	proto := determineProto()

	target, err := parseTarget(pflag.Args()[0])
	if err != nil {
		return nil
	}

	srcIPs, err := parseSources(*sources, target, proto)
	if err != nil {
		return nil
	}

	srcPortList := []uint16{0, 0}
	dstPortList := []uint16{0, 0}
	if proto == "udp" || proto == "tcp" {
		srcPortList, err = calcPorts(*srcPorts)
		if err != nil {
			return nil
		}
		dstPortList, err = calcPorts(*dstPorts)
		if err != nil {
			return nil
		}
	}

	if *length < 20 {
		*length = 20
	}

	if *maxTTL <= 0 || *timeout <= 0 || *flowInterval <= 0 || *printInterval <= 0 || *count <= 0 {
		return nil
	}

	return &core.TraceTask{
		DstIP:           target,
		SrcIPs:          srcIPs,
		DstPortMin:      dstPortList[0],
		DstPortMax:      dstPortList[1],
		SrcPortMin:      srcPortList[0],
		SrcPortMax:      srcPortList[1],
		Proto:           proto,
		IPVer:           4,
		MaxTTL:          *maxTTL,
		MsgLen:          *length,
		TOS:             *tos,
		TimeoutMs:       *timeout,
		FlowIntervalMs:  *flowInterval,
		PrintIntervalMs: *printInterval,
		Count:           *count,
	}
}

func determineProto() string {
	proto := "icmp"
	if *udp {
		proto = "udp"
	} else if *tcp {
		proto = "tcp"
	}

	return proto
}

func calcPorts(ports string) ([]uint16, error) {

	var minPort uint16
	var maxPort uint16
	portRange := strings.Split(ports, "-")
	if len(portRange) == 1 {
		port, err := strconv.ParseUint(portRange[0], 10, 16)
		if err != nil {
			fmt.Println("ntrace: unsupported port value")
			return nil, err
		}
		minPort = uint16(port)
		maxPort = uint16(port)
	} else if len(portRange) == 2 {
		low, err := strconv.ParseUint(portRange[0], 10, 16)
		if err != nil {
			fmt.Println("ntrace: unsupported port value")
			return nil, err
		}
		high, _ := strconv.ParseUint(portRange[1], 10, 16)
		if err != nil {
			fmt.Println("ntrace: unsupported port value")
			return nil, err
		}
		if low > high {
			fmt.Println("ntrace: unsupported port value")
			return nil, err
		}
		minPort = uint16(low)
		maxPort = uint16(high)
	} else {
		fmt.Println("ntrace: unsupported port value")
		return nil, errors.New("ntrace: unsupported port value")
	}

	parsedPorts := []uint16{minPort, maxPort}
	return parsedPorts, nil
}

func parseTarget(target string) (string, error) {

	var targets []string
	trimmed := strings.TrimSpace(target)
	if trimmed != "" {
		parsedTargets, err := nets.ParseIPAddr(trimmed)
		if err != nil {
			fmt.Println("ntrace: unsupported target")
			return "", err
		}
		targets = append(targets, parsedTargets...)
	}

	if len(targets) == 0 {
		fmt.Println("ntrace: please provide target")
		return "", errors.New("ntrace: please provide target")
	}

	return targets[0], nil
}

func parseSources(sources string, target string, proto string) ([]string, error) {
	var srcIPs []string

	if sources != "" {
		segs := strings.Split(sources, ",")
		for _, seg := range segs {
			trimmed := strings.TrimSpace(seg)
			if trimmed != "" {
				parsedSources, err := nets.ParseIPAddr(trimmed)
				if err != nil {
					fmt.Println("ntrace: unsupported source")
					return nil, err
				}
				srcIPs = append(srcIPs, parsedSources...)
			}
		}
	}

	if len(srcIPs) == 0 {
		srcIP, err := nets.LookupLocalAddrByRoute(target)
		if err != nil {
			return nil, err
		}
		srcIPs = append(srcIPs, srcIP.To4().String())
	}

	return srcIPs, nil
}

func helpFunc() bool {
	pflag.Usage = func() {
		fmt.Printf("Usage of %s:\n", os.Args[0])
		fmt.Printf("%s [-options] target\n", os.Args[0])
		pflag.PrintDefaults()
	}

	pflag.Parse()
	if *help {
		pflag.Usage()
		return true
	}

	if pflag.NArg() == 0 {
		pflag.Usage()
		return true
	}

	return false
}
